#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : ESIM.py
# @Author: Richard Chiming Xu
# @Date  : 2022/1/20
# @Desc  :

import re
import jieba
import pandas as pd
import numpy as np
from gensim.models import Word2Vec
import pickle
from tqdm import tqdm


class Config():
    # 数据加载部分
    dataset = 'paws-x'
    build_vocab = False  # 是否重构词典
    build_wv = False  # 是否重构词向量
    load_with_words = True  # 加载带分词的数据
    vocab_path = 'vocab'  # 路径
    wv_model_path = 'wv_model'  # 词向量名
    seq_len = 27  # 句子长度
    # 模型部分
    model_type = 'RNN'
    embed_dim = 100  # 词向量或embedding维度
    vocab_size = 100  # 词典大小，读取词典后会更新
    update_embed = True  # embedding是否随网络训练更新
    load_model = False  # 是否加载已有模型预测
    save_model = True  # 是否保存训练好的模型
    # 训练部分
    device = 'cpu'
    learning_rate = 5e-4
    batch_size = 128  # batch大小
    epochs = 25  # 训练次数
    print_loss = 100  # 打印loss次数
    dropout = 0.2   # dropout比例
    hidden_size = 256  # rnn的大小
    num_layers = 2  # rnn的层数
    concat_layers = True  # 是否拼接
    num_labels = 2  # 分类数


# 读取数据
def get_stopwords():
    stop_words = []
    with open('dataprocess/baidu_stopwords.txt', 'r', encoding='utf-8') as f:
        for line in f.readlines():
            stop_words.append(line.replace('\n', ''))
    return stop_words


# jieba分词
def cut(content, stop_words, config):
    # 去除符号
    content = re.sub("[\s+\.\!\/_,$%^*(+\"\']+|[+——！，。？、~@#￥%……&*（）]", "", content)

    result = []

    seg_list = jieba.lcut(content, cut_all=True)
    for i in seg_list:
        if i not in stop_words:
            result.append(i)

    if len(result) < config.seq_len:  # 小于规定长度，填充
        new_result = ['PAD' for i in range(config.seq_len)]
        new_result[:len(result)] = result
        return new_result
    else:
        return result[:config.seq_len]


# 保存pickle
def dump_pickle(obj, file_path):
    with open(file_path, 'wb') as f:
        pickle.dump(obj, f, pickle.HIGHEST_PROTOCOL)


# 加载pickle
def load_pickle(file_path):
    with open(file_path, 'rb') as f:
        return pickle.load(f)


# 分词并去除停用词
def data_anaysis(df, stop_words, config):
    # 分词
    df['words_a'] = df['text_a'].apply(lambda x: cut(x, stop_words, config))
    df['words_b'] = df['text_b'].apply(lambda x: cut(x, stop_words, config))
    return df


# 将中文词转换为词典数字
def word2num(content, vocab):
    result = []
    for word in content:
        result.append(vocab[word])
    return result


# 构建词典
def build_vocab(train, dev, test, config, file_name):
    word_index = 1
    vocab_dict = {
        'PAD': 0  # 默认pad为0
    }
    dataset = [train, test, dev]
    for data in tqdm(dataset, desc='构建词典'):
        for i in range(len(data)):
            row = data.iloc[i]
            words_a = row['words_a']
            words_b = row['words_b']
            for word in words_a:
                if word not in vocab_dict.keys():
                    vocab_dict[word] = word_index
                    word_index += 1
            for word in words_b:
                if word not in vocab_dict.keys():
                    vocab_dict[word] = word_index
                    word_index += 1
    # 保存词典
    dump_pickle(vocab_dict, config.vocab_path + '/vocab_' + file_name + '.pkl')
    return vocab_dict


def data_reader(file_name, config):
    print('加载数据...')
    if config.load_with_words is True:
        train = load_pickle('data/' + file_name + '_train.pkl')
        dev = load_pickle('data/' + file_name + '_dev.pkl')
        test = load_pickle('data/' + file_name + '_test.pkl')
    else:
        # 获取停用词
        stop_words = get_stopwords()
        train = pd.read_csv('data/' + file_name + '/train.tsv', sep='\t', names=['text_a', 'text_b', 'label'])
        dev = pd.read_csv('data/' + file_name + '/dev.tsv', sep='\t', names=['text_a', 'text_b', 'label'])
        test = pd.read_csv('data/' + file_name + '/test.tsv', sep='\t', names=['text_a', 'text_b'])

        if len(set(train['label'])) > 2:
            train = train[train['label'].isin(['0', '1'])]
            train['label'] = train['label'].astype('int')

        if len(set(train['label'])) > 2:
            dev = dev[dev['label'].isin(['0', '1'])]
            dev['label'] = dev['label'].astype('int')

        test['label'] = -1
        train = train.dropna()
        dev = dev.dropna()

        # 分词并去除停用词
        print('分词并去除停用词...')
        train = data_anaysis(train, stop_words, config)
        dev = data_anaysis(dev, stop_words, config)
        test = data_anaysis(test, stop_words, config)
        # 保存分词后的数据
        dump_pickle(train, 'data/' + file_name + '_train.pkl')
        dump_pickle(train, 'data/' + file_name + '_dev.pkl')
        dump_pickle(test, 'data/' + file_name + '_test.pkl')

    # 统计词典
    print('加载词典...')
    if config.build_vocab is True:
        vocab = build_vocab(train, dev, test, config, file_name)
    else:
        vocab = load_pickle(config.vocab_path + '/vocab_' + file_name + '.pkl')

    # 将每个词转为数字
    print('将每个词转为词典数字...')
    train['word2num_a'] = train['words_a'].apply(lambda x: word2num(x, vocab))
    train['word2num_b'] = train['words_b'].apply(lambda x: word2num(x, vocab))
    dev['word2num_a'] = dev['words_a'].apply(lambda x: word2num(x, vocab))
    dev['word2num_b'] = dev['words_b'].apply(lambda x: word2num(x, vocab))
    test['word2num_a'] = test['words_a'].apply(lambda x: word2num(x, vocab))
    test['word2num_b'] = test['words_b'].apply(lambda x: word2num(x, vocab))

    return train, dev, test, vocab


# 加载词向量
def load_wv(train, dev, test, config, file_name):
    if config.build_wv is True:
        print('构建词向量...')
        # 若不存在词向量，则重新训练训练词向量
        datalist = [train, dev, test]

        context = []
        for data in datalist:
            for i in range(len(data)):
                row = data.iloc[i]
                context.append(row['word2num_a'])
                context.append(row['word2num_b'])

        wv_model = Word2Vec(sentences=context, vector_size=config.embed_dim, window=5, min_count=1, workers=4)
        wv_model.train(context, total_examples=1, epochs=1)
        wv_model.save(config.wv_model_path + '/wv_' + file_name)
    else:
        print('加载词向量...')
        wv_model = Word2Vec.load(config.wv_model_path + '/wv_' + file_name)

    return wv_model


from torch.utils.data import DataLoader, Dataset


# 自定义数据集
class SimDataset(Dataset):
    def __init__(self, df):
        super(SimDataset, self).__init__()
        self.text_a = df['word2num_a']
        self.text_b = df['word2num_b']
        self.label = df['label']
        self.len = len(df)

    def __getitem__(self, idx):
        vector_a = np.array(self.text_a.iloc[idx], dtype='int64')
        vector_b = np.array(self.text_b.iloc[idx], dtype='int64')
        label = np.array(self.label.iloc[idx]).astype("int64")

        return {'vector_a': vector_a,
                'vector_b': vector_b,
                'label': label}

    def __len__(self):
        return self.len


from torch import nn
from torch.nn import functional as F
import torch
import torch.nn as nn


class RNNDropout(nn.Dropout):
    # 将词向量 某些维度 清0
    def forward(self, sequences_batch):  # (B, L, D)
        # 创建相同的全1张量 (B, D)
        ones = sequences_batch.data.new_ones(sequences_batch.shape[0], sequences_batch.shape[-1])

        # 创建随机mask (B, D)
        dropout_mask = nn.functional.dropout(ones, self.p, self.training, inplace=False)

        # dropout原数据 (B, L, D)， 这里需要给mask加一个维度
        return dropout_mask.unsqueeze(1) * sequences_batch


# 自定义RNN
class StackedBRNN(nn.Module):
    def __init__(self, input_size, hidden_size, num_layers,
                 dropout_rate=0, dropout_output=False, rnn_type=nn.LSTM,
                 concat_layers=False):
        super().__init__()
        # 获取参数
        self.dropout_output = dropout_output
        self.dropout_rate = dropout_rate
        self.num_layers = num_layers
        self.concat_layers = concat_layers  # 使用最后一层结果或叠加结果
        self.rnns = nn.ModuleList()
        # 遍历设计的RNN层数，使用Modulelist堆叠
        for i in range(num_layers):
            # 如果不是第一层，把lstm的两个hidden_size作为输入
            if i != 0:
                input_size = 2 * hidden_size
            self.rnns.append(rnn_type(input_size, hidden_size, num_layers=1, bidirectional=True))

    def forward(self, x):  # (B, L, D)
        # 转化成RNN能接收的维度
        x = x.transpose(0, 1)  # (L, B, D)

        # 用于记录不同层的RNN结果，初始是x
        outputs = [x]
        for i in range(self.num_layers):
            rnn_input = outputs[-1]

            # dropout
            if self.dropout_rate > 0:
                rnn_input = F.dropout(rnn_input, p=self.dropout_rate, training=self.training)

            # 取上一层的RNN结果传入当前层的RNN
            rnn_output = self.rnns[i](rnn_input)[0]  # 只获取output,无需使用h_n,c_n
            # 添加结果
            outputs.append(rnn_output)

        if self.concat_layers:  # 如果使用拼接作为结果
            # 这里0是X输入，所以只需要1开始取各层RNN的结果
            output = torch.cat(outputs[1:], 2)  # (L, B, D)
        else:  # 如果使用最后一层RNN作为结果
            output = outputs[-1]  # (L, B, D)

        # 还原维度
        output = output.transpose(0, 1)  # (B, L, D)

        # dropout
        if self.dropout_output and self.dropout_rate > 0:
            output = F.dropout(output, p=self.dropout_rate, training=self.training)  # (B, L, D)

        # 进行 transpose之后，tensor在内存中不连续， contiguous将output内存连续
        return output.contiguous()


class BidirectionalAttention(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, v1, v1_mask, v2, v2_mask):
        '''
            v1 (B, L, H)
            v1_mask (B, L)
            v2 (B, R, H)
            v2_mask (B, R)
        '''
        # v2:a v1:b

        # 1.计算矩阵相似度 v1@v2
        similarity_matrix = v1.bmm(v2.transpose(2, 1).contiguous())  # (B, L, R)

        # 2.计算attention时没有必要计算pad=0, 要进行mask操作 3.进行softmax
        # 将similarity_matrix v1中pad对应的权重给mask
        # v1_mask (B, L) 加一维到第三维度成 (B, L, unsqueeze)

        v2_v1_attn = F.softmax(
            similarity_matrix.masked_fill(
                v1_mask.unsqueeze(2), -1e7), dim=1)  # (B, L, R)

        # 将similarity_matrix v2中pad对应的权重给mask
        # 21_mask (B, R) 加一维到第三维度成 (B, unsqueeze, R)
        v1_v2_attn = F.softmax(
            similarity_matrix.masked_fill(
                v2_mask.unsqueeze(1), -1e7), dim=2)  # (B, L, R)

        # 4.计算attention
        # 句子b 对a的影响
        # attented_v1 (B, L, R) @ (B, R, H)
        attented_v1 = v1_v2_attn.bmm(v2)  # (B, L, H)

        # 句子b 对a的影响
        # v2_v1_attn (B, L, R) -> (B, R, L) @(B, L, H)
        attented_v2 = v2_v1_attn.transpose(1, 2).bmm(v1)  # (B, R, H)

        # attented_v1 将v1对应的pad填充为0
        # attented_v2 将v2对应的pad填充为0
        attented_v1.masked_fill(v1_mask.unsqueeze(2), 0)
        attented_v2.masked_fill(v2_mask.unsqueeze(2), 0)
        return attented_v1, attented_v2


class ESIM(nn.Module):

    def __init__(self, wv_mode, config: Config):
        super(ESIM, self).__init__()

        # -----------------------  encoding  ---------------------#
        word_vectors = torch.randn([config.vocab_size, config.embed_dim])
        for i in range(0, config.vocab_size):
            word_vectors[i, :] = torch.from_numpy(wv_mode.wv[i])
        # 创建embedding层
        self.embedding = nn.Embedding.from_pretrained(word_vectors, freeze=config.update_embed)  # (32, 27, 100)
        if config.update_embed is False:
            self.embedding.weight.requires_grad = False

        # 创建rnn的dropout
        self.rnn_dropout = RNNDropout(config.dropout)
        rnn_size = config.hidden_size

        if config.concat_layers is True:
            rnn_size //= config.num_layers
            config.hidden_size = rnn_size // 2 *2 *2 # 第一个*2是双向 第二个*2是cat拼接

        self.input_encoding = StackedBRNN(input_size=config.embed_dim,
                                          hidden_size=rnn_size // 2,
                                          num_layers=config.num_layers,
                                          rnn_type=nn.LSTM,
                                          concat_layers=config.concat_layers)

        # -----------------------  encoding  ---------------------#

        # -----------------------  注意力  ---------------------#
        self.attention = BidirectionalAttention()
        # -----------------------  注意力  ---------------------#

        # -----------------------  组合层  ---------------------#
        self.projection = nn.Sequential(
            nn.Linear(4 * config.hidden_size, config.hidden_size),
            nn.ReLU()
        )

        self.composition = StackedBRNN(input_size=config.hidden_size,
                                       hidden_size=rnn_size // 2,
                                       num_layers=config.num_layers,
                                       rnn_type=nn.LSTM,
                                       concat_layers=config.concat_layers)
        # -----------------------  组合层  ---------------------#

        # -----------------------  推理层 ---------------------#
        self.classification = nn.Sequential(
            nn.Dropout(p=config.dropout),
            nn.Linear(4 * config.hidden_size, config.hidden_size),
            nn.Tanh(),
            nn.Dropout(p=config.dropout))

        self.out = nn.Linear(config.hidden_size, config.num_labels)
        # -----------------------  推理层 ---------------------#

    def forward(self, words_a, words_b):

        '''
            维度说明表
                B: batch_size
                L: 句子a的长度
                R: 句子b的长度
                D: embedding长度
                H: hidden长度
        '''

        # 读取数据
        query = words_a  # (B, L)
        doc = words_b  # (B, R)

        # -----------------------  encoding  ---------------------#
        # 获取mask,判断query,doc种每个数是不是0
        # 是0则表示该位置是PAD
        # 是1则表示该位置不是PAD
        # query: [2,3,4,5,0,0,0] -> query_mask: [0,0,0,0,1,1,1]
        query_mask = (query == 0)  # (B, L)
        doc_mask = (query == 0)  # (B, R)

        # 转换词向量
        query = self.embedding(query)  # (B, L, D)
        doc = self.embedding(doc)  # (B, R, D)

        # dropout,随机对输出清零
        query = self.rnn_dropout(query)  # (B, L, D)
        doc = self.rnn_dropout(doc)  # (B, R, D)

        # 使用ESIM叠加的双向RNN 进行编码
        query = self.input_encoding(query)  # (B, L, H)
        doc = self.input_encoding(doc)  # (B, R, H)
        # -----------------------  encoding  ---------------------#

        # -----------------------  注意力  ---------------------#
        '''
            1. 计算两个句子的矩阵相似度
            2. 把PAD填充去掉，因为计算attention时先进行mask操作
            3. 进行softmax
            3. 计算attention
        '''
        attended_query, attended_doc = self.attention(query, query_mask, doc, doc_mask)
        # -----------------------  注意力  ---------------------#

        # -----------------------  拼接层  ---------------------#
        # 得到拼接embedding和attention得到加强信息版query和doc, 对应论文中的m
        enhanced_query = torch.cat([query, attended_query, query - attended_query, query * attended_query],
                                   dim=-1)  # (B, L, 4*H)
        enhanced_doc = torch.cat([doc, attended_doc, query - attended_doc, query * attended_doc], dim=-1)  # (B, R, 4*H)
        # -----------------------  拼接层  ---------------------#

        # -----------------------  组合层  ---------------------#
        # 推理拼接后的张量, 对应论文中的F(m)
        projected_query = self.projection(enhanced_query)  # (B, L, H)
        projected_doc = self.projection(enhanced_doc)  # (B, R, H)

        # 使用双向RNN
        query = self.composition(projected_query)  # (B, L, H)
        doc = self.composition(projected_doc)  # (B, R, H)
        # -----------------------  组合层  ---------------------#

        # -----------------------  池化层 ---------------------#
        '''
            1. 平均池化
            2. 最大池化
            3. 拼接 4个结果张量
        '''
        # 由于部分句子被pad，使用平均池化会不准，所以需要反推mask，然后求句子长度
        # 0的位置为pad
        reverse_query_mask = 1. - query_mask.float()  # (B, L)
        reverse_doc_mask = 1. - doc_mask.float()  # (B, R)

        # 平均池化
        query_avg = torch.sum(query * reverse_query_mask.unsqueeze(2), dim=1) / (
                torch.sum(reverse_query_mask, dim=1, keepdim=True) + 1e-8)  # (B, L, H)
        doc_avg = torch.sum(doc * reverse_doc_mask.unsqueeze(2), dim=1) / (
                torch.sum(reverse_doc_mask, dim=1, keepdim=True) + 1e-8)  # (B, R, H)

        # 防止取出pad（也许部分值是负数，小于0）
        query = query.masked_fill(query_mask.unsqueeze(2), -1e7)
        doc = doc.masked_fill(doc_mask.unsqueeze(2), -1e7)
        # 最大池化
        query_max, _ = query.max(dim=1)  # (B, L, H)
        doc_max, _ = doc.max(dim=1)  # (B, R, H)

        # 拼接
        X = torch.cat([query_avg, query_max, doc_avg, doc_max], dim=-1)
        # -----------------------  池化层 ---------------------#

        # -----------------------  推理层 ---------------------#
        X = self.classification(X)
        output = self.out(X)
        # -----------------------  推理层 ---------------------#

        return output


from torch.optim import AdamW
from torch.optim.lr_scheduler import ExponentialLR
from sklearn.metrics import f1_score, accuracy_score


def predict(config, model, test_df):
    predict_labels = []
    # 创建dataloader
    test_dataset = SimDataset(test_df)
    test_dataloader = DataLoader(test_dataset, batch_size=config.batch_size)
    model.eval()
    with torch.no_grad():
        for mini_batch in test_dataloader:
            mini_batch = {item: value.to(config.device) for item, value in mini_batch.items()}
            # 获取数据
            text_a = mini_batch['vector_a']
            text_b = mini_batch['vector_b']
            y_pred = model(text_a, text_b)
            _, indices = torch.max(y_pred, dim=1)
            predict_labels += indices
    return torch.tensor(predict_labels).numpy()


# 校验
def evaluation(config, model, val_dataloader, loss_fn):
    model.eval()
    preds = []
    labels = []
    val_loss = 0.
    # val_iterator = tqdm(val_dataloader, desc='Evaluation', total=len(val_dataloader))
    with torch.no_grad():
        for mini_batch in val_dataloader:
            mini_batch = {item: value.to(config.device) for item, value in mini_batch.items()}
            # 获取数据
            text_a = mini_batch['vector_a']
            text_b = mini_batch['vector_b']
            label = mini_batch['label']
            labels += label
            y_pred = model(text_a, text_b)

            loss = loss_fn(y_pred, label)

            val_loss += loss.item()
            # 返回逻辑值最大的位置，要么0，要么1
            _, indices = torch.max(y_pred, dim=1)
            preds += indices

    avg_val_loss = val_loss / len(val_dataloader)
    labels = torch.tensor(labels).numpy()
    preds = torch.tensor(preds).numpy()
    f1 = f1_score(labels, preds, average='macro')
    # -----------new ----------------#
    acc = accuracy_score(labels, preds)
    # -----------new ----------------#
    return avg_val_loss, f1, acc


#  训练数据
def train(train_df, dev_df, wv_model, config):
    # 创建数据集
    train_dataset = SimDataset(train_df)
    dev_dataset = SimDataset(dev_df)
    train_dataloader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True, drop_last=True)
    dev_dataloader = DataLoader(dev_dataset, batch_size=config.batch_size)

    # 创建模型
    model = ESIM(wv_model, config)

    model.to(config.device)
    # 定义优化器
    opt = AdamW(lr=config.learning_rate, params=model.parameters())
    # scheduler = ExponentialLR(opt, gamma=0.9)
    # 定义损失函数
    loss_fn = nn.CrossEntropyLoss()

    # 遍历epoch，开始训练
    for epoch in range(config.epochs):
        model.train()
        for iter_id, mini_batch in enumerate(train_dataloader):
            mini_batch = {item: value.to(config.device) for item, value in mini_batch.items()}
            # 获取数据
            text_a = mini_batch['vector_a']
            text_b = mini_batch['vector_b']
            label = mini_batch['label']
            # 训练
            y_pred = model(text_a, text_b)
            # 计算损失
            loss = loss_fn(y_pred, label)
            # 打印loss
            _, indices = torch.max(y_pred, dim=1)
            correct = torch.sum(indices == label)

            avg_loss = torch.mean(loss)
            if iter_id % config.print_loss == 0:
                print('epoch: {}, iter: {}, loss is: {}, acc is: {}'.format(epoch, iter_id, avg_loss,
                                                                            correct.item() * 1.0 / len(text_a)))
            # 反向传播
            avg_loss.backward()
            # 最小化loss,更新参数
            opt.step()
            # scheduler.step()
            # 清除梯度
            model.zero_grad()

        # 完成1个epoch，验证
        avg_val_loss, f1, acc = evaluation(config, model, dev_dataloader, loss_fn)
        print('-' * 50)
        print('epoch: {}, val_loss: {}, val_f1: {}, val_acc: {}'.format(epoch, avg_val_loss, f1, acc))
        print('-' * 50)

    return model


data_list = ['bq_corpus', 'lcqmc', 'paws-x']

for data in data_list:
    # 配置文件
    conf = Config()
    conf.device = 'cuda' if torch.cuda.is_available() else 'cpu'
    conf.dataset = data
    # 读取数据
    train_df, dev_df, test_df, vocab = data_reader(conf.dataset, conf)

    if conf.load_model is False:
        # 加载测向量
        wv_model = load_wv(train_df, dev_df, test_df, conf, conf.dataset)
        # 更新词典大小
        conf.vocab_size = len(vocab)
        # 训练模型
        model = train(train_df, dev_df, wv_model, conf)
        if conf.save_model is True:
            torch.save(model, 'model/' + conf.dataset + '.pkl')
    else:
        model = torch.load('model/' + conf.dataset + '.pkl')

    # 推理模型
    predict_labels = predict(conf, model, test_df)
    # 保存结果
    test_df['index'] = test_df.index
    test_df['prediction'] = predict_labels
    test_df.to_csv('result/' + conf.dataset + '.tsv', index=False, columns=['index', 'prediction'], sep='\t')
